import os
import cv2
import time
import xml.dom.minidom
import numpy as np
from lxml import etree
import random
from tqdm import tqdm as tqdm

# 把图片里被标为delete的目标，都改成像素值为127区域
anno_path = r'/Users/edy/Data/yq1-0-show/yq1_0_show_data_mix_dp/dp_mix_split_ten_dp/second_train/raw_data/Annotations'
img_path = r'/Users/edy/Data/yq1-0-show/yq1_0_show_data_mix_dp/dp_mix_split_ten_dp/second_train/raw_data/JPEGImages'

out_anno_dir = r'/Users/edy/Data/yq1-0-show/yq1_0_show_data_mix_dp/dp_mix_split_ten_dp/second_train/Annotations'
out_img_dir = r'/Users/edy/Data/yq1-0-show/yq1_0_show_data_mix_dp/dp_mix_split_ten_dp/second_train/JPEGImages'
if os.path.exists(out_anno_dir) is False:
    os.makedirs(out_anno_dir)
if os.path.exists(out_img_dir) is False:
    os.makedirs(out_img_dir)
remove_classes = ['delete']


def recursive_parse_xml_to_dict(xml):
    if not xml:
        return {xml.tag: xml.text}
    result = {}
    for child in xml:
        child_result = recursive_parse_xml_to_dict(child)
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


anno_list = list(filter(lambda x: x.endswith('.xml'), os.listdir(anno_path)))
for i in tqdm(anno_list):
    xmlfile = os.path.join(anno_path, i)
    with open(xmlfile, 'rb') as f:
        xml_str = f.read()
    img = cv2.imread(os.path.join(img_path, i.replace('.xml', '.jpg')))
    xml_dict = recursive_parse_xml_to_dict(etree.fromstring(xml_str))
    if 'object' in xml_dict['annotation']:
        for j in range(len(xml_dict['annotation']['object']) - 1, -1, -1):
            if xml_dict['annotation']['object'][j]["name"] in remove_classes:
                info = xml_dict['annotation']['object'][j]
                xml_dict['annotation']['object'].pop(j)
                empty = np.ones((int(info['bndbox']['ymax']) - int(info['bndbox']['ymin']),
                                 int(info['bndbox']['xmax']) - int(info['bndbox']['xmin']), 3), dtype="uint8")
                empty[:] = (127, 127, 127)
                img[int(info['bndbox']['ymin']):int(info['bndbox']['ymax']),
                int(info['bndbox']['xmin']):int(info['bndbox']['xmax'])] = empty
                # img[]
                continue
        cv2.imwrite(os.path.join(out_img_dir, i.replace('.xml', '.jpg')), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
        # cv2.imshow('',img)
        # cv2.waitKey()
    doc = xml.dom.minidom.Document()
    # owner
    annotation = doc.createElement('annotation')
    doc.appendChild(annotation)
    # owner
    folder = doc.createElement('folder')
    annotation.appendChild(folder)
    folder_txt = doc.createTextNode("JPEGImages")
    folder.appendChild(folder_txt)

    filename = doc.createElement('filename')
    annotation.appendChild(filename)
    filename_txt = doc.createTextNode(xml_dict['annotation']['filename'])
    filename.appendChild(filename_txt)

    path = doc.createElement('path')
    annotation.appendChild(path)
    if 'path' not in xml_dict['annotation']:
        xml_dict['annotation']['path'] = xml_dict['annotation']['filename']
    path_txt = doc.createTextNode(xml_dict['annotation']['path'])
    path.appendChild(path_txt)
    # ones#
    source = doc.createElement('source')
    annotation.appendChild(source)

    database = doc.createElement('database')
    source.appendChild(database)
    database_txt = doc.createTextNode("Unknown")
    database.appendChild(database_txt)

    size = doc.createElement('size')
    annotation.appendChild(size)

    width = doc.createElement('width')
    size.appendChild(width)
    width_txt = doc.createTextNode(xml_dict['annotation']['size']['width'])
    width.appendChild(width_txt)

    height = doc.createElement('height')
    size.appendChild(height)
    height_txt = doc.createTextNode(xml_dict['annotation']['size']['height'])
    height.appendChild(height_txt)

    depth = doc.createElement('depth')
    size.appendChild(depth)
    depth_txt = doc.createTextNode("3")
    depth.appendChild(depth_txt)
    # twoe#
    segmented = doc.createElement('segmented')
    annotation.appendChild(segmented)
    segmented_txt = doc.createTextNode("0")
    segmented.appendChild(segmented_txt)

    if 'object' not in xml_dict['annotation']:
        continue
    for j in range(0, int(len(xml_dict['annotation']['object']))):
        info = xml_dict['annotation']['object'][j]
        # print(objbuds)
        # threes#
        object_new = doc.createElement("object")
        annotation.appendChild(object_new)

        name = doc.createElement('name')
        object_new.appendChild(name)

        name_txt = doc.createTextNode(info['name'])
        name.appendChild(name_txt)

        pose = doc.createElement('pose')
        object_new.appendChild(pose)
        pose_txt = doc.createTextNode("Unspecified")
        pose.appendChild(pose_txt)

        truncated = doc.createElement('truncated')
        object_new.appendChild(truncated)
        truncated_txt = doc.createTextNode("0")
        truncated.appendChild(truncated_txt)

        difficult = doc.createElement('difficult')
        object_new.appendChild(difficult)
        difficult_txt = doc.createTextNode("0")
        difficult.appendChild(difficult_txt)
        # threes-1#
        bndbox = doc.createElement('bndbox')
        object_new.appendChild(bndbox)

        xmin = doc.createElement('xmin')
        bndbox.appendChild(xmin)
        xmin_txt = doc.createTextNode(info['bndbox']['xmin'])
        xmin.appendChild(xmin_txt)

        ymin = doc.createElement('ymin')
        bndbox.appendChild(ymin)
        ymin_txt = doc.createTextNode(info['bndbox']['ymin'])
        ymin.appendChild(ymin_txt)

        xmax = doc.createElement('xmax')
        bndbox.appendChild(xmax)
        xmax_txt = doc.createTextNode(info['bndbox']['xmax'])
        xmax.appendChild(xmax_txt)

        ymax = doc.createElement('ymax')
        bndbox.appendChild(ymax)
        ymax_txt = doc.createTextNode(info['bndbox']['ymax'])
        ymax.appendChild(ymax_txt)

    with open(os.path.join(out_anno_dir, i), "w", encoding='utf-8') as f:
        f.write(doc.toprettyxml(indent='\t').replace('<?xml version="1.0" ?>\n', ''))
